from langgraph.graph import StateGraph, END

from test.langgraph_base.ticket_node import intent_recognition_node, category_node, assign_node
from test.langgraph_base.ticket_state import TicketState

# 1. 初始化图（绑定状态类型）
workflow = StateGraph(TicketState)

# 2. 添加节点（name为节点标识，func为节点函数）
workflow.add_node("intent_recognition", intent_recognition_node)
workflow.add_node("category", category_node)
workflow.add_node("assign", assign_node)

# 定义条件函数：根据意图决定下一个节点
# 如果是售后/咨询/投诉，继续到分类节点
# 如果是其他，直接到结束节点
def decide_after_intent(state):
    """根据意图识别结果决定下一个节点"""
    # 获取当前状态的意图（需要处理dict或TicketState两种情况）
    intent = state.get('intent', '') if isinstance(state, dict) else getattr(state, 'intent', '')

    # 条件判断
    # if intent in ['售后', '咨询', '投诉']:
    if intent in ['售后', '咨询']:
        print(f"条件路由：意图'{intent}'，转到分类节点")
        return "category"
    else:
        print(f"条件路由：意图'{intent}'，直接结束流程")
        return END


# 定义条件函数：根据分类决定是否需要分配

def decide_after_category(state):
    """根据工单分类决定是否需要分配人员"""
    # 获取当前状态的分类
    category = state.get('category', '') if isinstance(state, dict) else getattr(state, 'category', '')

    # 条件判断
    if category in ['产品故障', '退款申请', '功能咨询', '价格咨询']:
        print(f"条件路由：分类'{category}'，转到分配节点")
        return "assign"
    else:
        print(f"条件路由：分类'{category}'，直接结束流程")
        return END



# 3. 定义边（流转规则）
# 起始节点 → 意图识别节点
workflow.set_entry_point("intent_recognition")

# 使用条件函数定义边（有条件流转）
# 意图识别 → 条件判断函数 → category或END
workflow.add_conditional_edges(
    "intent_recognition",  # 源节点
    decide_after_intent    # 条件函数
)

# 分类节点 → 条件判断函数 → assign或END
workflow.add_conditional_edges(
    "category",          # 源节点
    decide_after_category  # 条件函数
)

# 分配节点 → 结束（无条件流转）
workflow.add_edge("assign", END)

# 4. 编译图（生成可执行对象）
app = workflow.compile()


if __name__ == '__main__':
    print("===== 工单处理系统测试 =====")
    
    # 测试场景1：退款申请工单
    print("\n[测试场景1] 退款申请工单")
    # 创建初始状态
    initial_state = TicketState(ticket_id="TICKET-001", user_query="我想退款，产品不太符合我的需求")
    print(f"初始状态: {initial_state.model_dump()}")
    
    # 执行完整流程
    final_state = app.invoke(initial_state)
    
    # 输出结果
    print(f"\n处理完成!")
    print(f"final_state: {final_state}")
    # 检查final_state类型并正确访问（使用字典方式访问）
    print(f"工单ID: {final_state.get('ticket_id', 'N/A')}")
    print(f"用户请求: {final_state.get('user_query', 'N/A')}")
    print(f"识别意图: {final_state.get('intent', 'N/A')}")
    print(f"工单分类: {final_state.get('category', 'N/A')}")
    print(f"分配人员: {final_state.get('assigned_to', 'N/A')}")
    
    print(f"\n执行历史:")
    for step in final_state.get('history', []):
        print(f"  - {step.get('node', 'N/A')}: {step.get('result', 'N/A')} (时间: {step.get('timestamp', 'N/A')})")
    
    print("\n===== 测试完成 =====")
    print("工单处理流程已成功执行。")
    print("LangGraph 流程图验证通过，正确处理了返回的字典格式结果。")
    
    # # 测试场景2：产品故障工单
    # print("\n[测试场景2] 产品故障工单")
    # initial_state = TicketState(ticket_id="TICKET-002", user_query="我的设备坏了，无法开机")
    # final_state = app.invoke(initial_state)
    # print(f"工单ID: {final_state.ticket_id}")
    # print(f"识别意图: {final_state.intent}")
    # print(f"工单分类: {final_state.category}")
    # print(f"分配人员: {final_state.assigned_to}")
    #
    # # 测试场景3：功能咨询工单
    # print("\n[测试场景3] 功能咨询工单")
    # initial_state = TicketState(ticket_id="TICKET-003", user_query="请问这个功能怎么用？")
    # final_state = app.invoke(initial_state)
    # print(f"工单ID: {final_state.ticket_id}")
    # print(f"识别意图: {final_state.intent}")
    # print(f"工单分类: {final_state.category}")
    # print(f"分配人员: {final_state.assigned_to}")
    #
    # # 测试场景4：价格咨询工单
    # print("\n[测试场景4] 价格咨询工单")
    # initial_state = TicketState(ticket_id="TICKET-004", user_query="这个产品价格是多少？有优惠吗？")
    # final_state = app.invoke(initial_state)
    # print(f"工单ID: {final_state.ticket_id}")
    # print(f"识别意图: {final_state.intent}")
    # print(f"工单分类: {final_state.category}")
    # print(f"分配人员: {final_state.assigned_to}")
    #
    # # 测试场景5：其他类型工单
    # print("\n[测试场景5] 其他类型工单")
    # initial_state = TicketState(ticket_id="TICKET-005", user_query="我想提个建议")
    # final_state = app.invoke(initial_state)
    # print(f"工单ID: {final_state.ticket_id}")
    # print(f"识别意图: {final_state.intent}")
    # print(f"工单分类: {final_state.category}")
    # print(f"分配人员: {final_state.assigned_to}")
    #
    # print("\n===== 测试完成 =====")
    # print("所有工单处理流程已成功执行。")
    # print("LangGraph 流程图验证通过，节点间正确流转。")

